#!/usr/bin/env python3

# Copyright © 2025 Wenze Wei
#
# This file is part of Pisces L1.
#
# Licensed under the Creative Commons Attribution-NonCommercial 4.0 International License (CC BY-NC 4.0).
# You may not use this file except in compliance with the License.
# Commercial use is strictly prohibited.
# You may obtain a copy of the License at
#
#     https://creativecommons.org/licenses/by-nc/4.0/
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import torch

def save_ckpt(model, optimizer, epoch, path):
    """
    Save the model checkpoint to the specified path.

    Args:
        model (torch.nn.Module): The PyTorch model to be saved.
        optimizer (torch.optim.Optimizer): The optimizer associated with the model.
        epoch (int): The current training epoch number.
        path (str): The path where the checkpoint will be saved.

    Returns:
        None: This function does not return any value.
    """
    # Create the directory if it doesn't exist
    os.makedirs(os.path.dirname(path), exist_ok=True)
    # Save the model state, optimizer state, and epoch number to the specified path
    torch.save({
        'model': model.state_dict(), 
        'optimizer': optimizer.state_dict(), 
        'epoch': epoch
    }, path)


def load_ckpt(path, model, optimizer):
    """
    Load the model checkpoint from the specified path.

    Args:
        path (str): The path where the checkpoint is located.
        model (torch.nn.Module): The PyTorch model to load the state into.
        optimizer (torch.optim.Optimizer): The optimizer to load the state into.

    Returns:
        int: The epoch number stored in the checkpoint.
    """
    # Load the checkpoint data to CPU
    ckpt = torch.load(path, map_location='cpu')
    # Load the model state from the checkpoint
    model.load_state_dict(ckpt['model'])
    # Load the optimizer state from the checkpoint
    optimizer.load_state_dict(ckpt['optimizer'])
    return ckpt['epoch']